"""
Author        yongfa
Date          2022-07-07 14:21:46
LastEditTime  2022-10-17 10:07:10
LastEditors   yongfa
Description   文件流下载
"""

from io import BytesIO

import pandas as pd
from flask import send_file, Response


class Downloader:

    def download_file(self, columns: list = [], data: list = [], file_name: str = 'test'):
        """文件下载

        Args:
            columns (list, optional): 文件行标题. Defaults to [].
            data (list, optional): 文件数据. Defaults to [].
            file_name (str, optional): 导出文件名称. Defaults to 'test.csv'.

        Returns:
            _type_: _description_
        """

        # df_1 = pd.DataFrame(np.random.randint(0, 10, size=(100, 8)), columns=list('ABCDEFGH'))

        df = pd.DataFrame(data, columns=columns)

        # create an output stream
        output = BytesIO()
        writer = pd.ExcelWriter(output, engine='xlsxwriter')
        file_full_name = file_name.split('.')[0] + '.xlsx'

        # taken from the original question
        df.to_excel(writer, startrow=0, merge_cells=False, sheet_name="{}".format(file_name))
        workbook = writer.book
        worksheet = writer.sheets[file_name]
        format = workbook.add_format()
        format.set_bg_color('#eeeeee')
        worksheet.set_column(1, len(columns), 20)

        # the writer has done its job
        writer.close()

        # go back to the beginning of the stream
        output.seek(0)

        # finally return the file
        return send_file(output, attachment_filename=file_full_name, as_attachment=True)

    def download_file_stream(self, columns: list = [], data: list = [], file_name: str = 'test.csv'):
        def generate():
            # 标题
            yield ','.join(columns) + '\n'

            # 数据
            for row in data:
                line = []
                for col in row:
                    line.append(str(col))

                yield ','.join(line) + '\n'

        response = Response(generate(), mimetype="text/csv")
        response.headers["Content-disposition"] = 'attachment; filename={}'.format(file_name)
        return response
